from PIL import Image
import numpy as np
from albumentations.pytorch.transforms import ToTensorV2
import os
import sys
current_path = os.path.abspath(__file__)
current_dir = os.path.dirname(current_path)
parent_dir = os.path.dirname(current_dir)
sys.path.append(parent_dir)
from image_to_latex.lit_models import LitResNetTransformer
import time


def calculate_time(func):
    def wrapper(*args, **kwargs):
        start_time = time.time()
        result = func(*args, **kwargs)
        end_time = time.time()
        print(f"Function {func.__name__} took {end_time - start_time} seconds to run.")
        return result
    return wrapper


file = 'scripts/test.png'


@calculate_time
def inference(file):
    image = Image.open(file).convert('L')
    lit_model = LitResNetTransformer.load_from_checkpoint("/home/openkylin/image-to-latex/artifacts/model.pt")
    lit_model.freeze()
    transform = ToTensorV2()

    image_tensor = transform(image=np.array(image))["image"]  # type: ignore
    pred = lit_model.model.predict(image_tensor.unsqueeze(0).float())[0]  # type: ignore
    decoded = lit_model.tokenizer.decode(pred.tolist())  # type: ignore
    decoded_str = "".join(decoded)
    return decoded_str

# result = inference(file)
# print(result)


def resize_image(image_path, new_width, new_height):
    original_image = Image.open(image_path)
    new_image = Image.new("RGBA", (new_width, new_height), "white")

    # 计算粘贴位置
    paste_x = (new_width - original_image.width) // 2
    paste_y = (new_height - original_image.height) // 2

    new_image.paste(original_image, (paste_x, paste_y))

    # 保存调整后的图像
    new_image.save("resized_image.png")


# 指定图像路径和目标尺寸
image_path = "test.png"
new_width = 500
new_height = 200

# 调用函数进行图像尺寸调整
resize_image(image_path, new_width, new_height)